# import os
# import json
#
# import torch
# from PIL import Image
# from torchvision import transforms
# import matplotlib.pyplot as plt
#
# from vit_model import vit_base_patch16_224_in21k as create_model
#
#
# def main():
#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
#     data_transform = transforms.Compose(
#         [transforms.Resize(256),
#          transforms.CenterCrop(224),
#          transforms.ToTensor(),
#          transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
#
#     # load image
#     img_path = r"D:\ViT\dataset\train\dog\dog.20.jpg"
#     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
#     img = Image.open(img_path)
#     plt.imshow(img)
#     # [N, C, H, W]
#     img = data_transform(img)
#     # expand batch dimension
#     img = torch.unsqueeze(img, dim=0)
#
#     # read class_indict
#     json_path = './class_indices.json'
#     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
#
#     with open(json_path, "r") as f:
#         class_indict = json.load(f)
#
#     # create model
#     model = create_model(num_classes=2, has_logits=False).to(device)
#     # load model weights
#     model_weight_path = "./model/transformer_model2.pth"
#     model.load_state_dict(torch.load(model_weight_path, map_location=device))
#     model.eval()
#     with torch.no_grad():
#         # predict class
#         output = torch.squeeze(model(img.to(device))).cpu()
#         predict = torch.softmax(output, dim=0)
#         predict_cla = torch.argmax(predict).numpy()
#
#     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
#                                                  predict[predict_cla].numpy())
#     plt.title(print_res)
#     for i in range(len(predict)):
#         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
#                                                   predict[i].numpy()))
#     plt.show()
#
#
# if __name__ == '__main__':
#     main()
import os
import json

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

from vit_model import vit_base_patch16_224_in21k as create_model


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose(
        [transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
         transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

    # load image
    img_path = r"D:\ViT\dataset\test\cat\cat.104.jpg"
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)
    plt.imshow(img)
    # [N, C, H, W]
    img = data_transform(img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    # read class_indict
    json_path = './class_indices3.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)

    with open(json_path, "r") as f:
        class_indict = json.load(f)

    # create model
    model = create_model(num_classes=2, has_logits=False).to(device)
    # load model weights
    model_weight_path = "./model/transformer_model4.pth"
    state_dict = torch.load(model_weight_path, map_location=device)
    # 移除多余的键
    unwanted_keys = ["pre_logits.fc.weight", "pre_logits.fc.bias"]
    for key in unwanted_keys:
        if key in state_dict:
            del state_dict[key]

    model.load_state_dict(state_dict, strict=False)
    model.eval()
    with torch.no_grad():
        # predict class
        output = torch.squeeze(model(img.to(device))).cpu()
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()

    print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                 predict[predict_cla].numpy())
    plt.title(print_res)
    for i in range(len(predict)):
        print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
                                                  predict[i].numpy()))
    plt.show()


if __name__ == '__main__':
    main()

# import os
# import json
#
# import torch
# from PIL import Image
# from torchvision import transforms
# import matplotlib.pyplot as plt
#
# from vit_model import vit_base_patch16_224_in21k as create_model
#
#
# def main():
#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
#     data_transform = transforms.Compose(
#         [transforms.Resize(256),
#          transforms.CenterCrop(224),
#          transforms.ToTensor(),
#          transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
#
#     # load image
#     img_path = r"D:\ViT\mydata\train\20\2a758feb7c4e_76.jpg"
#     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
#     img = Image.open(img_path)
#     plt.imshow(img)
#     # [N, C, H, W]
#     img = data_transform(img)
#     # expand batch dimension
#     img = torch.unsqueeze(img, dim=0)
#
#     # read class_indict
#     json_path = './class_indices1.json'
#     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
#
#     with open(json_path, "r") as f:
#         class_indict = json.load(f)
#
#     # create model
#     model = create_model(num_classes=5, has_logits=False).to(device)
#     # load model weights
#     model_weight_path = "./model/transformer_model3.pth"
#     state_dict = torch.load(model_weight_path, map_location=device)
#     # 移除多余的键
#     unwanted_keys = ["pre_logits.fc.weight", "pre_logits.fc.bias"]
#     for key in unwanted_keys:
#         if key in state_dict:
#             del state_dict[key]
#
#     model.load_state_dict(state_dict, strict=False)
#     model.eval()
#     with torch.no_grad():
#         # predict class
#         output = torch.squeeze(model(img.to(device))).cpu()
#         predict = torch.softmax(output, dim=0)
#         predict_cla = torch.argmax(predict).numpy()
#
#     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
#                                                  predict[predict_cla].numpy())
#     plt.title(print_res)
#     for i in range(len(predict)):
#         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
#                                                   predict[i].numpy()))
#     plt.show()
#
#
# if __name__ == '__main__':
#     main()
